/*
 *  carAgent.cpp
 *
 *
 *  Created by Adam White on 24/02/06.
 *  Modified by Ali Nouri in April 2010 to be a general agent and not specific for mountain car.
 *  And also to take parameter values from ParamReader
 */
#include "QLearningAgent.h"
#include "ParamReader.h"

//variables -----------------
int dimension;								//dimension of the problem
int numberOfActions;                             // number of actions


double* tileWidths;							//the tile widths for each dimension of the state variable
int* f;                 // sets of features, one set per tiling
double* Q;
int** F;					//sets of features, one set per action


#define MEMORY_SIZE 10000                        // number of parameters to theta, memory size
#define MAX_NONZERO_TRACES 1000
int nonzero_traces[MAX_NONZERO_TRACES];
int num_nonzero_traces = 0;
int nonzero_traces_inverse[MEMORY_SIZE];
float minimum_trace = 0.01;
int nonzero_traces_g[MAX_NONZERO_TRACES];
int num_nonzero_traces_g = 0;
int nonzero_traces_inverse_g[MEMORY_SIZE];
float minimum_trace_g = 0.01;
// Global RL variables:

double theta[MEMORY_SIZE];                        // modifyable parameter vector, aka memory, weights
double omega[MEMORY_SIZE];                        // modifyable parameter vector, aka memory, weights

double e[MEMORY_SIZE];                            // eligibility traces

double g[MEMORY_SIZE];                            // eligibility traces
float F_hat[MEMORY_SIZE];

// Standard RL parameters:
#define gamma 1.0                      // discount-rate parameters

action_t oldAction;						//action selected on previous time step
double V;						//action selected on current time step
float * state_vars;

taskspec_t *ts;					//task spec


//--------------



float actor_policy_bottom()
{
	float bottom = 0.0;
	for (int b = 0; b < numberOfActions; b++)
		bottom += expf(Q[b]);
	return bottom;
}
float actor_policy(int a,int debug = 0)
{
	return expf(Q[a])/actor_policy_bottom();
}

int choose_action()
{

    float* probs = new float[numberOfActions];
    float tot = 0.0;
    float rand = (float)drand48();
	float bottom = actor_policy_bottom();
    for (int b = 0; b < numberOfActions; b++)
    {
        tot += expf(Q[b])/bottom;
        probs[b] = tot;
        if (rand < tot)
		{
			delete[] probs;
            return b;
		}
    }
//    printf("ERROR! probabilities don't add up to 1! %f,%f %f %f\n",rand,probs[0],probs[1],probs[2]);
	printf("Error: probabilities dont add up to 1\n\n\n");

    for (int b = 0; b < 3; b++)
    {
        actor_policy(b,1);
    }
    exit(1);

}

void agent_init(const char* _ts)
{
	//we have to make sure the load function is called in the main.cpp file, otherwise we need to uncomment the next line
//	ParamReader::load("param.in");
	ParamReader::print();



	ts=(taskspec_t*)malloc(sizeof(taskspec_t));


	for (int i=0; i<MEMORY_SIZE; i++){
		theta[i]= 0.0;                     // clear memory at start of each run
		omega[i]= 0.0;                   // clear memory at start of each run
		e[i] = 0.0;                            // clear all traces
	}

	//we dont use the epsilon?
//	ParamReader::AGNT_EPSILON = 0.0;

	allocateRLStruct(&oldAction, 1, 0, 0);
	int decode_result = decode_taskspec( ts, _ts);
	if(decode_result!=0){
		printf("Could not decode task spec, code: %d for task spec: %s\n",decode_result,_ts);
		exit(1);
	}

	//compute tile width
	dimension = ts->num_double_observations;
	numberOfActions = ts->int_actions[0].max - ts->int_actions[0].min +1;


	tileWidths = new double[dimension];
	Q = new double[numberOfActions];
	f = new int[ParamReader::AGNT_NUMBER_OF_TILINGS];
	F = Allocate2DArray<int>(numberOfActions, ParamReader::AGNT_NUMBER_OF_TILINGS);

	for(int i=0; i< dimension; i++)
		tileWidths[i] = (ts->double_observations[i].max - ts->double_observations[i].min)/ParamReader::AGNT_TILE_RESOLUTION;

	state_vars = new float[dimension];

}



const action_t* agent_start(const observation_t* o)
{
    DecayTraces(0.0);                                            // clear all traces

    LoadF(o->doubleArray);                                                     // compute features
    LoadQ();                                                     // compute action values
	SetV();

	if(drand48() <= ParamReader::AGNT_EPSILON)
		oldAction.intArray[0] = (int) (drand48() * (numberOfActions));
	else
		oldAction.intArray[0] = choose_action();                                       // pick argmax action

	return &oldAction;


}

const action_t* agent_step(double r, const observation_t* o)
{
  /*
  //warning:
  if (o[2] > 0)
    return 0;
  else
    return 2;

  */

	DecayTraces(gamma*ParamReader::AGNT_LAMBDA);                               // let traces fall
        for (int a=0; a<numberOfActions; a++)                        // optionally clear other traces
            if (a != oldAction.intArray[0])
                for (int j=0; j<ParamReader::AGNT_NUMBER_OF_TILINGS; j++) ClearTrace(F[a][j]);
	for (int j=0; j<ParamReader::AGNT_NUMBER_OF_TILINGS; j++) SetTrace(f[j],1.0);
	for (int j=0; j<ParamReader::AGNT_NUMBER_OF_TILINGS; j++) SetTrace_g(F[oldAction.intArray[0]][j],1.0);

	double delta = r - V;

	LoadF(o->doubleArray);
	LoadQ();

	SetV();

	if(drand48() <= ParamReader::AGNT_EPSILON)
		oldAction.intArray[0] = (int) (drand48() * (numberOfActions));
	else
		oldAction.intArray[0] = choose_action();

	delta += gamma*V;
	float temp = (ParamReader::AGNT_ALPHA/ParamReader::AGNT_NUMBER_OF_TILINGS)*delta;
	float temp2 = (ParamReader::AGNT_BETA/ParamReader::AGNT_NUMBER_OF_TILINGS)*delta;

	for (int i=0; i<num_nonzero_traces; i++)                 // update theta (learn)
	{
		int f = nonzero_traces[i];
		theta[f] += temp * e[f];
	}
	for (int i=0; i<num_nonzero_traces_g; i++)
	{
		int f = nonzero_traces_g[i];
		omega[f] += temp2 * g[f];
	}

	return &oldAction;
}

void agent_end(double r)
{
	DecayTraces(gamma*ParamReader::AGNT_LAMBDA);                               // let traces fall
        for (int a=0; a<numberOfActions; a++)                        // optionally clear other traces
            if (a != oldAction.intArray[0])
                for (int j=0; j<ParamReader::AGNT_NUMBER_OF_TILINGS; j++) ClearTrace(F[a][j]);
	for (int j=0; j<ParamReader::AGNT_NUMBER_OF_TILINGS; j++) SetTrace(f[j],1.0);
	for (int j=0; j<ParamReader::AGNT_NUMBER_OF_TILINGS; j++) SetTrace_g(F[oldAction.intArray[0]][j],1.0);

	double delta = r - V;

	float temp = (ParamReader::AGNT_ALPHA/ParamReader::AGNT_NUMBER_OF_TILINGS)*delta;
	float temp2 = (ParamReader::AGNT_BETA/ParamReader::AGNT_NUMBER_OF_TILINGS)*delta;

	for (int i=0; i<num_nonzero_traces; i++)                 // update theta (learn)
	{
		int f = nonzero_traces[i];
		theta[f] += temp * e[f];
	}
	for (int i=0; i<num_nonzero_traces_g; i++)
	{
		int f = nonzero_traces_g[i];
		omega[f] += temp2 * g[f];
	}

}


void agent_cleanup()
{
	delete[] state_vars;
	delete[] tileWidths;
	delete[] Q;
	delete[] f;
	Free2DArray(F);
	clearRLStruct(&oldAction);
}

void SetV()
{
	V = 0.0;
	for (int j = 0; j < ParamReader::AGNT_NUMBER_OF_TILINGS; j++)
		V += theta[f[j]];
}

void LoadQ()
// Compute all the action values from current F and theta
  { for (int a=0; a<numberOfActions; a++)
      { Q[a] = 0;
        for (int j=0; j<ParamReader::AGNT_NUMBER_OF_TILINGS; j++) Q[a] += omega[F[a][j]];
		if (Q[a] > 50.0) Q[a] = 50.0;
		if (Q[a] < -100.0) Q[a] = -100.0;
		}}

void LoadQ(int a)
// Compute an action value from current F and theta
  { Q[a] = 0;
    for (int j=0; j<ParamReader::AGNT_NUMBER_OF_TILINGS; j++) Q[a] += omega[F[a][j]];}

void LoadF(double* o)
// Compute feature sets for current car state
  {

	for(int i=0; i< dimension; i++)
		state_vars[i] = o[i] / tileWidths[i];

	tiles(&f[0],ParamReader::AGNT_NUMBER_OF_TILINGS,MEMORY_SIZE,state_vars,dimension);
    for (int a=0; a<numberOfActions; a++)
      tiles(&F[a][0],ParamReader::AGNT_NUMBER_OF_TILINGS,MEMORY_SIZE,state_vars,dimension,a);}


void SetTrace(int f, float new_trace_value)
// Set the trace for feature f to the given value, which must be positive
  { if (e[f] >= minimum_trace) e[f] = new_trace_value;         // trace already exists
    else { while (num_nonzero_traces >= MAX_NONZERO_TRACES) IncreaseMinTrace(); // ensure room for new trace
           e[f] = new_trace_value;
           nonzero_traces[num_nonzero_traces] = f;
           nonzero_traces_inverse[f] = num_nonzero_traces;
           num_nonzero_traces++;}}

void SetTrace_g(int f, float new_trace_value)
{
	if (g[f] >= minimum_trace_g) g[f] = new_trace_value;
	else { while (num_nonzero_traces_g >= MAX_NONZERO_TRACES) IncreaseMinTrace_g();
			g[f] = new_trace_value;
			nonzero_traces_g[num_nonzero_traces_g] = f;
			nonzero_traces_inverse_g[f] = num_nonzero_traces_g;
			num_nonzero_traces_g++;}
}

void ClearTrace(int f)
// Clear any trace for feature f
{
	if (!(e[f]==0.0))
        ClearExistentTrace(f,nonzero_traces_inverse[f]);
}

void ClearTrace_g(int f)
{
	if (!(g[f]==0.0))
		ClearExistentTrace(f,nonzero_traces_inverse_g[f]);
}


void ClearExistentTrace(int f, int loc)
// Clear the trace for feature f at location loc in the list of nonzero traces
  { e[f] = 0.0;
    num_nonzero_traces--;
    nonzero_traces[loc] = nonzero_traces[num_nonzero_traces];
    nonzero_traces_inverse[nonzero_traces[loc]] = loc;}

void ClearExistentTrace_g(int f, int loc)
{
	g[f] = 0.0;
	num_nonzero_traces_g--;
	nonzero_traces_g[loc] = nonzero_traces_g[num_nonzero_traces_g];
	nonzero_traces_inverse_g[nonzero_traces_g[loc]] = loc;
}

void DecayTraces(float decay_rate)
// Decays all the (nonzero) traces by decay_rate, removing those below minimum_trace
{
	for (int loc=num_nonzero_traces-1; loc>=0; loc--)      // necessary to loop downwards
    {
		int f = nonzero_traces[loc];
        e[f] *= decay_rate;
        if (e[f] < minimum_trace) ClearExistentTrace(f,loc);
	}
	for (int loc=num_nonzero_traces_g-1; loc >=0; loc--)
	{
		int f = nonzero_traces_g[loc];
		g[f] *= decay_rate;
		if (g[f] < minimum_trace_g) ClearExistentTrace_g(f,loc);
	}
}

void IncreaseMinTrace()
// Try to make room for more traces by incrementing minimum_trace by 10%,
// culling any traces that fall below the new minimum
  { minimum_trace += 0.1 * minimum_trace;
    //cout << "Changing minimum_trace to " << minimum_trace << endl;
    for (int loc=num_nonzero_traces-1; loc>=0; loc--)      // necessary to loop downwards
      { int f = nonzero_traces[loc];
        if (e[f] < minimum_trace) ClearExistentTrace(f,loc);};}

void IncreaseMinTrace_g()
{
	minimum_trace_g += 0.1 * minimum_trace_g;
	for (int loc = num_nonzero_traces_g-1; loc >= 0; loc--)
	{
		int f = nonzero_traces_g[loc];
		if (g[f] < minimum_trace_g) ClearExistentTrace_g(f,loc);
	}
}

const char* agent_message(const char* inMessage) {
	if(strcmp(inMessage,"what is your name?")==0)
		return "my name is Continuous QLearning Agent!";

	return "I don't know how to respond to your message";
}


template < typename T >
T **Allocate2DArray( int nRows, int nCols)
{
    T **ppi;
    T *pool;
    T *curPtr;
    //(step 1) allocate memory for array of elements of column

    ppi = new T*[nRows];

    //(step 2) allocate memory for array of elements of each row
    pool = new T [nRows * nCols];

    // Now point the pointers in the right place
    curPtr = pool;
    for( int i = 0; i < nRows; i++)
    {
        *(ppi + i) = curPtr;
         curPtr += nCols;
    }
    return ppi;
}

template < typename T >
void Free2DArray(T** Array)
{
    delete [] *Array;
    delete [] Array;
}
